import os
import pandas as pd
from loguru import logger
import multiprocessing as mp
import numpy as np

num_worker = 16
num_gpus = 4

def train_kpi(args):
    index, kpi_id, cuda_to_use = args

    logger.info(f'------------------> Training and detecting {kpi_id} ({index+1}/{len(kpi_id_list)})')
    cmd = f'CUDA_VISIBLE_DEVICES={cuda_to_use} python3 donut_train.py --kpi-id {kpi_id}'
    logger.info(cmd)
    os.system(cmd)

if __name__ == '__main__':
    kpi_id_list = sorted(set(pd.read_csv('data/train.csv', engine='c')['KPI ID'].values))
    logger.info(f"{len(kpi_id_list)} kpis in total.")
    cuda_to_use = np.random.randint(0, num_gpus, len(kpi_id_list)) if num_gpus > 0 else -np.ones(len(kpi_id_list), dtype=np.int64)
    print(cuda_to_use)
    index = list(range(len(kpi_id_list)))

    with mp.Pool(processes=num_worker) as p:
        p.map(train_kpi, iterable=zip(index, kpi_id_list, cuda_to_use))
